package com.gitee.c0dehvb.reactive.sample.r2dbc;

import com.gitee.c0dehvb.reactive.sample.r2dbc.connection.ShardingConnectionFactory;
import com.gitee.c0dehvb.reactive.sample.r2dbc.metadata.R2dbcSchemaMetaDataLoader;
import com.gitee.c0dehvb.reactive.sample.r2dbc.metadata.R2dbcSchemaMetaDataLoaderFactory;
import com.google.common.collect.Lists;
import io.r2dbc.spi.ConnectionFactory;
import org.apache.shardingsphere.api.config.sharding.ShardingRuleConfiguration;
import org.apache.shardingsphere.core.rule.ShardingRule;
import org.apache.shardingsphere.core.yaml.config.sharding.YamlTableRuleConfiguration;
import org.apache.shardingsphere.core.yaml.swapper.ShardingRuleConfigurationYamlSwapper;
import org.apache.shardingsphere.sharding.rewrite.context.ShardingSQLRewriteContextDecorator;
import org.apache.shardingsphere.sharding.route.engine.ShardingRouteDecorator;
import org.apache.shardingsphere.spi.database.type.DatabaseType;
import org.apache.shardingsphere.sql.parser.SQLParserEngine;
import org.apache.shardingsphere.sql.parser.binder.metadata.schema.SchemaMetaData;
import org.apache.shardingsphere.underlying.common.config.DatabaseAccessConfiguration;
import org.apache.shardingsphere.underlying.common.config.properties.ConfigurationProperties;
import org.apache.shardingsphere.underlying.common.database.type.DatabaseTypes;
import org.apache.shardingsphere.underlying.common.metadata.ShardingSphereMetaData;
import org.apache.shardingsphere.underlying.common.metadata.datasource.DataSourceMetas;
import org.apache.shardingsphere.underlying.rewrite.SQLRewriteEntry;
import org.apache.shardingsphere.underlying.route.DataNodeRouter;
import org.springframework.boot.autoconfigure.r2dbc.R2dbcProperties;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.r2dbc.core.DatabaseClient;
import org.springframework.util.StringUtils;

import javax.annotation.Resource;
import java.util.HashMap;
import java.util.Map;
import java.util.Properties;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
 * 分片R2dbc数据源配置类
 * @author LiYilin
 * @since 2022/6/14 10:51 AM
 **/
@Configuration
@EnableConfigurationProperties({SpringBootShardingRuleConfigurationProperties.class})
public class ShardingR2dbcConfiguration {
    private static final String DEFAULT_DATA_NODE = "default";
    private static final String ACTUAL_DATA_NODE_SPLITTER = ".";
    private static final Pattern WITH_SCHEMA_PATTERN = Pattern.compile("^\\w+\\.");

    @Resource
    private Properties properties;
    @Resource
    private R2dbcProperties r2dbcProperties;

    @Bean
    public DatabaseClient shardingDatabaseClient(ConnectionFactory connectionFactory,
                                                 ShardingRewriteEngine shardingRewriteEngine) {
        if (connectionFactory instanceof ShardingConnectionFactory) {
            return DatabaseClient.builder().connectionFactory(connectionFactory).build();
        } else {
            return DatabaseClient.builder().connectionFactory(new ShardingConnectionFactory(connectionFactory,
                                                                                            shardingRewriteEngine)).build();
        }
    }

    /**
     * 构建分表SQL重写引擎
     * @param connectionFactory
     * @param shardingRule
     * @param r2dbcProperties
     * @param properties
     * @return
     */
    @Bean
    public ShardingRewriteEngine shardingRewriteEngine(ConnectionFactory connectionFactory,
                                                       SpringBootShardingRuleConfigurationProperties shardingRule,
                                                       R2dbcProperties r2dbcProperties,
                                                       Properties properties) {
        ConfigurationProperties configurationProperties = new ConfigurationProperties(properties);

        // 构建分表规则
        // 目前仅支持单数据源，需要保证分表规则中的actual-data-nodes配置只有一种数据源
        String actualDataNode = getOnlyOneActualDataNode(shardingRule);
        setDefaultActualDataNode(shardingRule, actualDataNode);
        ShardingRuleConfiguration configuration = new ShardingRuleConfigurationYamlSwapper().swap(shardingRule);
        ShardingRule rule = new ShardingRule(configuration, Lists.newArrayList(actualDataNode));

        // 分析数据源类型
        DatabaseType databaseType = getDatabaseTypeByUrl(r2dbcProperties.getUrl());
        if (databaseType == null) {
            throw new IllegalArgumentException("不支持的数据库类型");
        }

        // 加载数据库表元数据
        SchemaMetaData schemaMetaData = loadSchemaMetaData(databaseType, rule, r2dbcProperties, connectionFactory);

        SQLParserEngine parserEngine = new SQLParserEngine(databaseType.getName());
        ShardingSphereMetaData shardingSphereMetaData =
                new ShardingSphereMetaData(createDataSourceMetas(r2dbcProperties), schemaMetaData);
        DataNodeRouter dataNodeRouter = new DataNodeRouter(shardingSphereMetaData, configurationProperties,
                                                           parserEngine);
        dataNodeRouter.registerDecorator(rule, new ShardingRouteDecorator());
        SQLRewriteEntry sqlRewriteEntry = new SQLRewriteEntry(schemaMetaData, configurationProperties);
        ShardingSQLRewriteContextDecorator sqlRewriteContextDecorator =
                new ShardingSQLRewriteContextDecorator();
        sqlRewriteEntry.registerDecorator(rule, sqlRewriteContextDecorator);
        return new ShardingRewriteEngine(schemaMetaData, parserEngine, dataNodeRouter, sqlRewriteEntry);
    }

    private String getOnlyOneActualDataNode(SpringBootShardingRuleConfigurationProperties shardingRule) {
        String actualDataNode = null;
        for (String table : shardingRule.getTables().keySet()) {
            final YamlTableRuleConfiguration yamlTableRuleConfiguration = shardingRule.getTables().get(table);
            String actualDataNodes = yamlTableRuleConfiguration.getActualDataNodes();
            if (StringUtils.hasText(actualDataNodes)) {
                Matcher matcher = WITH_SCHEMA_PATTERN.matcher(actualDataNodes);
                if (matcher.find()) {
                    String group = matcher.group(0);
                    String dataNode = group.substring(0, group.length() - 1);
                    if (actualDataNode == null) {
                        actualDataNode = dataNode;
                    } else {
                        throw new IllegalArgumentException("暂不支持多数据源，请保证每个actual-data-nodes配置的数据源是一致的");
                    }
                }
            }
        }
        return actualDataNode == null ? DEFAULT_DATA_NODE : actualDataNode;
    }

    private void setDefaultActualDataNode(SpringBootShardingRuleConfigurationProperties shardingRule,
                                          String actualDataNode) {
        shardingRule.getTables().forEach((table, yamlTableRuleConfiguration) -> {
            String actualDataNodes = yamlTableRuleConfiguration.getActualDataNodes();
            if (StringUtils.hasText(actualDataNodes)) {
                Matcher matcher = WITH_SCHEMA_PATTERN.matcher(actualDataNodes);
                if (!matcher.find()) {
                    yamlTableRuleConfiguration.setActualDataNodes(actualDataNode + ACTUAL_DATA_NODE_SPLITTER + actualDataNodes);
                } else {
                    yamlTableRuleConfiguration.setActualDataNodes(matcher.replaceFirst(actualDataNode + ACTUAL_DATA_NODE_SPLITTER));
                }
            }
        });
    }

    private SchemaMetaData loadSchemaMetaData(DatabaseType databaseType,
                                              ShardingRule rule,
                                              R2dbcProperties r2dbcProperties,
                                              ConnectionFactory connectionFactory) {
        R2dbcSchemaMetaDataLoader loader = R2dbcSchemaMetaDataLoaderFactory.buildByDatabaseType(databaseType);
        if (loader == null) {
            throw new IllegalArgumentException("不支持的数据库类型");
        }
        return loader.load(rule, r2dbcProperties.getPool().getMaxSize(), connectionFactory);
    }

    private DataSourceMetas createDataSourceMetas(R2dbcProperties r2dbcProperties) {
        Map<String, DatabaseAccessConfiguration> databaseAccessConfigurationMap = new HashMap<>();
        // DataSourceMetas需要通过jdbcUrl分析数据库类型，替换成jdbc开头适配一下，不会真建立连接
        String url = r2dbcProperties.getUrl().replace("r2dbcs", "jdbc")
                .replace("r2dbc", "jdbc");
        DatabaseType databaseType = DatabaseTypes.getDatabaseTypeByURL(url);
        DatabaseAccessConfiguration databaseAccessConfiguration =
                new DatabaseAccessConfiguration(url, r2dbcProperties.getUsername(),
                                                r2dbcProperties.getPassword());
        databaseAccessConfigurationMap.put(DEFAULT_DATA_NODE, databaseAccessConfiguration);
        return new DataSourceMetas(databaseType, databaseAccessConfigurationMap);
    }

    private DatabaseType getDatabaseTypeByUrl(String url) {
        url = r2dbcProperties.getUrl().replace("r2dbcs", "jdbc")
                .replace("r2dbc", "jdbc");
        return DatabaseTypes.getDatabaseTypeByURL(url);
    }
}
